import numpy as np
from scipy import signal
from conv import get_output_size
from scipy.ndimage import filters

if __name__ == '__main__':

    def get_son_array(array, row, col, shape):
        return array[:, row:row+shape[0], col:col+shape[1]]

    def conv_py(input_array, filter_array, step):
        if input_array.shape[0] != filter_array.shape[0]:
            print("conv deep error, input_array deep %s, filter deep %s" %(input_array.shape[0], filter_array.shape[0]))
            assert(False)
            return
        output_array_height = get_output_size(input_array.shape[1], filter_array.shape[1], step)
        output_array_width = get_output_size(input_array.shape[2], filter_array.shape[2], step)
        output_array = np.zeros((output_array_height, output_array_width))
        filter_shape = filter_array.shape[1:]
        #filter_shape = filter_array.shape
        for i in range(output_array_height):
            for j in range(output_array_width):
                output_array[i][j] = sum((get_son_array(input_array, i*step, j*step, filter_shape)*filter_array).flat)
        #output_array = np.array([(get_son_array(input_array, i*step, j*step, filter_array.shape[1:])*filter_array).sum() \
        #    for i in range(output_array_height) for j in range(output_array_width)]).reshape(output_array_height, output_array_width)
        return output_array

    x = np.array([  [1, 1, 1, 0, 0],
                    [0, -1, 1, 1, 0],
                    [0, 0, -1, 1, 1],
                    [0, 0, 1, 1, 0],
                    [0, 1, 1, 0, 0]])

    #x = np.array([[0, 1, 1, 0, 2],
    #      [2, 2, 2, 2, 1],
    #      [1, 0, 0, 2, 0],
    #      [0, 1, 1, 0, 0],
    #      [1, 2, 0, 0, 2]])

    w_k = np.array([[1, 1, 0],
                    [0, 1, 0],
                    [0, 0, 0]])


    f1 = signal.convolve2d(x, np.rot90(w_k, 2), 'valid')
    print(f1)
    print("---------")

    d = np.fromiter((x+1 for x in w_k.flat), dtype=w_k.dtype).reshape(w_k.shape)

    y = np.array((x, x, x))
    y_w_k = np.array(
        [[[1, 1, 0], [2, 4, -3], [8, 1, 0]],
         [[1, 0, 1], [-1, -5, 2], [1, 8, 10]],
         [[-1, -2, 3], [0, 1, 0], [2, 4, 7]]])
    f = signal.convolve(y, np.rot90(y_w_k, 2, (1, 2)), 'valid')[0]
    print(f)
    print(conv_py(y, y_w_k, 1))

    a = ((1, 2), (3, 4))  
    b = ('m', 'n')
    print("---------")
    #c = [l for l in x1 for x1, x2 in zip(a, b)]
    c = [ l for x1, x2 in zip(a, b) for l in x1]
    print(c)

    with open("train.out", "w") as fp:
        fp.write("test result")

    def padding(input_array, zp):
        input_d = input_array.shape[0]
        input_h = input_array.shape[1]
        input_w = input_array.shape[2]
        padding_array = np.zeros((input_d, input_h+2*zp, input_w+2*zp))
        padding_array[:, zp:zp+input_h, zp:zp+input_w] = input_array
        return padding_array

    input_array = np.array(
        [[[0, 1, 1, 0, 2],
          [2, 2, 2, 2, 1],
          [1, 0, 0, 2, 0],
          [0, 1, 1, 0, 0],
          [1, 2, 0, 0, 2]],
         [[1, 0, 2, 2, 0],
          [0, 0, 0, 2, 0],
          [1, 2, 1, 2, 1],
          [1, 0, 0, 0, 0],
          [1, 2, 1, 1, 1]],
         [[2, 1, 2, 0, 0],
          [1, 0, 0, 1, 0],
          [0, 2, 1, 0, 1],
          [0, 1, 2, 2, 2],
          [2, 1, 0, 0, 1]]])

    filter_w =  np.array(
      [[[-0.37019064,  0.33079862,  0.10745912],
        [-0.1418555 ,  0.47251358, -0.39377449],
        [-0.25936635,  0.73677354, -0.85959855]],
       [[-0.49996249, -0.73649416,  0.38663987],
        [ 0.47489245,  0.85442775,  0.27314919],
        [-0.98686727, -0.44811474,  0.06283342]],
       [[-0.29495747,  0.09487862, -0.19475832],
        [ 0.302171  ,  0.59140534, -0.31582818],
        [ 0.55775404, -0.16458845, -0.95216453]]])

   # print("========")
   # print(np.rot90(filter_w, 2, (2, 1)))
    print("========")
    print(signal.convolve(padding(input_array, 1), np.rot90(filter_w, 2, (2, 1)), 'valid'))
    #print(signal.convolve(padding(input_array, 1), filter_w, 'valid')[0])

    x1 = np.array(
         [[[0, 1, 1, 0, 2],
          [2, 2, 2, 2, 1],
          [1, 0, 0, 2, 0],
          [0, 1, 1, 0, 0],
          [1, 2, 0, 0, 2]]]
    )

    y1 = np.array(
       [[[-0.37019064,  0.33079862,  0.10745912],
        [-0.1418555 ,  0.47251358, -0.39377449],
        [-0.25936635,  0.73677354, -0.85959855]]])

    print("++++++++")
    print(signal.convolve2d(padding(x1, 1)[0], np.rot90(y1[0], 2), 'valid'))


    print("***********", y.shape, input_array.shape)
    f = signal.convolve(padding(y, 1), np.rot90(filter_w, 2, (1, 2)), 'valid')[0]
    print(f)
    print(conv_py(padding(y, 1), filter_w, 1))

    def swap_array(ary, d1, d2):
        tmp = np.copy(ary[d1])
        ary[d1] = ary[d2]
        ary[d2] = tmp


    print("***********", y.shape, input_array.shape)
    #input_array[2], input_array[0] = input_array[0], input_array[2]
    print(conv_py(padding(input_array, 1), filter_w, 1))
    
    rot_y_w_k = np.rot90(y_w_k, 2, (2, 1))
    #rot_y_w_k = np.roll(rot_y_w_k, 1, axis=0)
    f = signal.convolve(padding(input_array, 1), np.flip(np.rot90(filter_w, 2, (2, 1)), 0), 'valid')[0]
    print(f)
    

#    print("++++++++-----")
#    print(conv_py(padding(x, 1) , y, 1))
#    print("++++++++-----")
#    print(conv_py(padding(input_array, 1) , filter_w, 1))